package com.slipper.ai.vector;

import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.SearchRequest;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

public class Vector {
    private final EmbeddingModel embeddingModel;
    private final Map<String, VectorContent> store = new ConcurrentHashMap<>();

    public Vector(EmbeddingModel embeddingModel) {
        this.embeddingModel = embeddingModel;
    }

    public void add(String id, String text, Map<String, Object> metadata, float[] embedding) {
        VectorContent content = new VectorContent(id, text, metadata, embedding);
        this.store.put(id, content);
    }

    public List<Document> similaritySearch(SearchRequest request) {
        float[] userQueryEmbedding = this.getQueryEmbedding(request.getQuery());
        return this.store.values().stream()
                .map(content -> content.toDocument(EmbeddingMath.cosineSimilarity(userQueryEmbedding, content.getEmbedding())))
                .filter(content -> content.getScore() != null && content.getScore() >= request.getSimilarityThreshold())
                .sorted(Comparator.comparing(Document::getScore).reversed())
                .limit(request.getTopK())
                .toList();
    }

    private float[] getQueryEmbedding(String query) {
        return this.embeddingModel.embed(query);
    }

    public static final class EmbeddingMath {
        private EmbeddingMath() {
            throw new UnsupportedOperationException("This is a utility class and cannot be instantiated");
        }

        public static double cosineSimilarity(float[] vectorX, float[] vectorY) {
            if (vectorX != null && vectorY != null) {
                if (vectorX.length != vectorY.length) {
                    throw new IllegalArgumentException("Vectors lengths must be equal");
                } else {
                    float dotProduct = dotProduct(vectorX, vectorY);
                    float normX = norm(vectorX);
                    float normY = norm(vectorY);
                    if (normX != 0.0F && normY != 0.0F) {
                        return (double)dotProduct / (Math.sqrt(normX) * Math.sqrt(normY));
                    } else {
                        throw new IllegalArgumentException("Vectors cannot have zero norm");
                    }
                }
            } else {
                throw new RuntimeException("Vectors must not be null");
            }
        }

        public static float dotProduct(float[] vectorX, float[] vectorY) {
            if (vectorX.length != vectorY.length) {
                throw new IllegalArgumentException("Vectors lengths must be equal");
            } else {
                float result = 0.0F;

                for(int i = 0; i < vectorX.length; ++i) {
                    result += vectorX[i] * vectorY[i];
                }

                return result;
            }
        }

        public static float norm(float[] vector) {
            return dotProduct(vector, vector);
        }
    }
}
